'''
难度 中等
题目链接 https://leetcode.cn/problems/find-the-punishment-number-of-an-integer/description/
'''

class Solution:
    # 回溯
    def punishmentNumber(self, n: int) -> int:
        def dfs(s: str, p: int, sum: int, target: int) -> bool:
            # 当切割位置等于长度 直接比较结果是否相等
            if p == len(s):
                return sum == target

            sumNew = 0
            for i in range(p, len(s)):
                sumNew = sumNew * 10 + int(s[i])
                if sum + sumNew > target:
                    break
                
                if dfs(s, i + 1, sum + sumNew, target):
                    return True

            return False
        
        结果 = 0
        for i in range(1, n + 1):
            if dfs(str(i * i), 0, 0, i):
                结果 += i * i
                # print(i*i)

        return 结果
    

if __name__ == '__main__':
    s = Solution()
    参数 = [
        [10],
        [37],
    ]
    预期 = [
        182,
        1478,
    ]
    for k,v in enumerate(参数):
        r = s.punishmentNumber(v[0])
        if r != 预期[k]:
            print("第{}个case不通过 期望={} 结果={}\n".format(
                k+1,
                预期[k],
                r
            ))